1use crate::ext::io::*;
3use crate::ext::psb::*;
4use crate::scripts::base::*;
5use crate::types::*;
6use crate::utils::encoding::*;
7use crate::utils::img::*;
8use anyhow::Result;
9use emote_psb::PsbReader;
10use libtlg_rs::TlgColorType;
11use std::collections::HashMap;
12use std::io::Read;
13use std::path::{Path, PathBuf};
14use url::Url;
15
16#[derive(Debug)]
17pub struct DrefBuilder {}
19
20impl DrefBuilder {
21 pub fn new() -> Self {
23 Self {}
24 }
25}
26
27impl ScriptBuilder for DrefBuilder {
28 fn default_encoding(&self) -> Encoding {
29 Encoding::Cp932
30 }
31
32 fn build_script(
33 &self,
34 buf: Vec<u8>,
35 filename: &str,
36 encoding: Encoding,
37 _archive_encoding: Encoding,
38 config: &ExtraConfig,
39 archive: Option<&Box<dyn Script>>,
40 ) -> Result<Box<dyn Script>> {
41 Ok(Box::new(Dref::new(
42 buf, encoding, filename, config, archive,
43 )?))
44 }
45
46 fn extensions(&self) -> &'static [&'static str] {
47 &["dref"]
48 }
49
50 fn script_type(&self) -> &'static ScriptType {
51 &ScriptType::EmoteDref
52 }
53
54 fn is_image(&self) -> bool {
55 true
56 }
57}
58
59struct Dpak {
60 psb: VirtualPsbFixed,
61}
62
63struct OffsetData {
64 left: u32,
65 top: u32,
66}
67
68impl Dpak {
69 pub fn new<P: AsRef<Path>>(path: P) -> Result<Self> {
70 let f = std::fs::File::open(path)?;
71 let mut f = std::io::BufReader::new(f);
72 let psb = PsbReader::open_psb_v2(&mut f)?;
73 let psb = psb.to_psb_fixed();
74 Ok(Self { psb })
75 }
76
77 pub fn load_from_data(data: &[u8]) -> Result<Self> {
78 let psb = PsbReader::open_psb_v2(MemReaderRef::new(data))?;
79 let psb = psb.to_psb_fixed();
80 Ok(Self { psb })
81 }
82
83 pub fn load_image(&self, name: &str) -> Result<(ImageData, Option<OffsetData>)> {
84 let root = self.psb.root();
85 let rid = root[name]
86 .resource_id()
87 .ok_or_else(|| anyhow::anyhow!("Resource ID for image '{}' not found in DPAK", name))?
88 as usize;
89 if rid >= self.psb.resources().len() {
90 return Err(anyhow::anyhow!(
91 "Resource ID {} out of bounds for DPAK with {} resources",
92 rid,
93 self.psb.resources().len()
94 ));
95 }
96 let resource = &self.psb.resources()[rid];
97 Self::load_img(&resource)
98 }
99
100 fn load_img(data: &[u8]) -> Result<(ImageData, Option<OffsetData>)> {
101 if libtlg_rs::is_valid_tlg(data) {
102 Ok((Self::load_tlg(data)?, None))
103 } else {
104 Self::load_png(data)
105 }
106 }
107
108 fn load_tlg(data: &[u8]) -> Result<ImageData> {
109 let img = libtlg_rs::load_tlg(MemReaderRef::new(data))
110 .map_err(|e| anyhow::anyhow!("Failed to decode TLG image: {:?}", e))?;
111 let color = img.color;
112 let mut re = ImageData {
113 width: img.width as u32,
114 height: img.height as u32,
115 color_type: match img.color {
116 TlgColorType::Grayscale8 => ImageColorType::Grayscale,
117 TlgColorType::Bgr24 => ImageColorType::Bgr,
118 TlgColorType::Bgra32 => ImageColorType::Bgra,
119 },
120 data: img.data,
121 depth: 8,
122 };
123 if let Some(v) = img.tags.get(&Vec::from(b"mode")) {
124 if v == b"alpha" && color == TlgColorType::Bgr24 {
125 convert_bgr_to_bgra(&mut re)?;
126 }
127 }
128 Ok(re)
129 }
130
131 fn load_png(data: &[u8]) -> Result<(ImageData, Option<OffsetData>)> {
132 let mut img = load_png(MemReaderRef::new(&data))?;
133 match img.color_type {
134 ImageColorType::Rgb => {
135 convert_rgb_to_rgba(&mut img)?;
136 }
137 _ => {}
138 }
139 Ok((
140 img,
141 Self::try_read_offset_from_png(MemReaderRef::new(&data))?,
142 ))
143 }
144
145 fn try_read_offset_from_png(mut data: MemReaderRef) -> Result<Option<OffsetData>> {
146 data.pos = 8; data.pos += 8; data.pos += 17; loop {
150 let chunk_size = data.read_u32_be()?;
151 let mut chunk_type = [0u8; 4];
152 data.read_exact(&mut chunk_type)?;
153 if &chunk_type == b"IDAT" || &chunk_type == b"IEND" {
154 break;
155 }
156 if &chunk_type == b"oFFs" {
157 let x = data.read_u32_be()?;
158 let y = data.read_u32_be()?;
159 if data.read_u8()? == 0 {
160 return Ok(Some(OffsetData { left: x, top: y }));
161 }
162 }
163 data.pos += chunk_size as usize + 4; }
165 Ok(None)
166 }
167}
168
169#[derive(Default)]
170struct DpakLoader {
171 map: HashMap<String, Dpak>,
172}
173
174impl DpakLoader {
175 pub fn load_image(
176 &mut self,
177 dir: &Path,
178 dpak: &str,
179 filename: &str,
180 ) -> Result<(ImageData, Option<OffsetData>)> {
181 let dpak = match self.map.get(dpak) {
182 Some(d) => d,
183 None => {
184 let mut path = dir.join(dpak);
185 path = crate::utils::files::get_ignorecase_path(&path)?;
186 let ndpak = Dpak::new(&path)?;
187 self.map.insert(dpak.to_string(), ndpak);
188 self.map.get(dpak).unwrap()
189 }
190 };
191 dpak.load_image(filename)
192 }
193
194 pub fn load_archives(&mut self, in_archives: &HashMap<String, Vec<u8>>) -> Result<()> {
195 for (name, data) in in_archives.iter() {
196 if !self.map.contains_key(name) {
197 let dpak = Dpak::load_from_data(data)?;
198 self.map.insert(name.clone(), dpak);
199 }
200 }
201 Ok(())
202 }
203}
204
205pub struct Dref {
207 urls: Vec<Url>,
208 dir: PathBuf,
209 in_archives: HashMap<String, Vec<u8>>,
210}
211
212impl std::fmt::Debug for Dref {
213 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
214 f.debug_struct("Dref")
215 .field("urls", &self.urls)
216 .field("dir", &self.dir)
217 .finish()
218 }
219}
220
221impl Dref {
222 pub fn new(
230 buf: Vec<u8>,
231 encoding: Encoding,
232 filename: &str,
233 _config: &ExtraConfig,
234 archive: Option<&Box<dyn Script>>,
235 ) -> Result<Self> {
236 let text = decode_with_bom_detect(encoding, &buf, true)?.0;
237 let mut urls = Vec::new();
238 for text in text.lines() {
239 let text = text.trim();
240 if text.is_empty() {
241 continue;
242 }
243 urls.push(Url::parse(text)?);
244 }
245 let path = Path::new(filename);
246 let dir = if let Some(parent) = path.parent() {
247 parent.to_path_buf()
248 } else {
249 PathBuf::from(".")
250 };
251 if urls.is_empty() {
252 return Err(anyhow::anyhow!("No URLs found in DREF file: {}", filename));
253 }
254 for u in urls.iter() {
255 if u.scheme() != "psb" {
256 return Err(anyhow::anyhow!(
257 "Invalid URL scheme in DREF file: {} (expected 'psb')",
258 u
259 ));
260 }
261 }
262 let mut in_archives = HashMap::new();
263 if let Some(archive) = archive {
264 if archive.is_archive() {
265 for url in urls.iter() {
266 let filename = url.domain().ok_or(anyhow::anyhow!(
267 "Invalid URL in DREF file: {} (missing domain)",
268 url
269 ))?;
270 if let Ok(mut content) = archive.open_file_by_name(filename, true) {
271 in_archives.insert(filename.to_string(), content.data()?);
272 }
273 }
274 }
275 }
276 Ok(Self {
277 urls,
278 dir,
279 in_archives,
280 })
281 }
282}
283
284impl Script for Dref {
285 fn default_output_script_type(&self) -> OutputScriptType {
286 OutputScriptType::Json
287 }
288
289 fn default_format_type(&self) -> FormatOptions {
290 FormatOptions::None
291 }
292
293 fn is_image(&self) -> bool {
294 true
295 }
296
297 fn export_image(&self) -> Result<ImageData> {
298 let mut loader = DpakLoader::default();
299 loader.load_archives(&self.in_archives)?;
300 let base_url = &self.urls[0];
301 let dpak = base_url.domain().ok_or(anyhow::anyhow!(
302 "Invalid URL in DREF file: {} (missing domain)",
303 base_url
304 ))?;
305 let (mut base_img, base_offset) =
306 loader.load_image(&self.dir, dpak, base_url.path().trim_start_matches("/"))?;
307 if let Some(o) = base_offset {
308 eprintln!("WARN: Base image offset: left={}, top={}", o.left, o.top);
309 crate::COUNTER.inc_warning();
310 }
311 for url in &self.urls[1..] {
312 let dpak = url.domain().ok_or(anyhow::anyhow!(
313 "Invalid URL in DREF file: {} (missing domain)",
314 url
315 ))?;
316 let (mut img, img_offset) =
317 loader.load_image(&self.dir, dpak, url.path().trim_start_matches("/"))?;
318 let (top, left) = match img_offset {
319 Some(o) => (o.top, o.left),
320 None => (0, 0),
321 };
322 if base_img.color_type != img.color_type {
323 if base_img.color_type == ImageColorType::Rgba
324 && img.color_type == ImageColorType::Rgb
325 {
326 convert_rgb_to_rgba(&mut img)?;
327 } else if base_img.color_type == ImageColorType::Bgra
328 && img.color_type == ImageColorType::Bgr
329 {
330 convert_bgr_to_bgra(&mut img)?;
331 } else if base_img.color_type == ImageColorType::Rgba
332 && img.color_type == ImageColorType::Bgra
333 {
334 convert_bgra_to_rgba(&mut img)?;
335 } else if base_img.color_type == ImageColorType::Bgra
336 && img.color_type == ImageColorType::Rgba
337 {
338 convert_rgba_to_bgra(&mut img)?;
339 }
340 }
341 draw_on_img_with_opacity(&mut base_img, &img, left, top, 0xff)?;
342 }
343 Ok(base_img)
344 }
345}